use std::path::Path;
use windows::Win32::Foundation::{ERROR_NO_MORE_ITEMS, ERROR_SUCCESS};
use windows::Win32::System::Registry::{
    HKEY, HKEY_LOCAL_MACHINE, KEY_READ, REG_SZ, REG_VALUE_TYPE, RegEnumKeyExW, RegOpenKeyExW,
    RegQueryValueExW,
};
use windows::core::HSTRING;

use crate::classes::*;
use crate::enums::*;
use crate::ntdll::*;

pub struct RegistryManager;

impl RegistryManager {
    pub fn new() -> Self {
        Self
    }

    fn option_value_to_option_state(value: OptionValue, force: bool, option: OptionValue) -> u64 {
        let num = match value {
            OptionValue::On => 1,
            OptionValue::Off => 2,
            _ => 0,
        };

        let mut option_state = num as u64;

        if force {
            option_state |= RtlImageMitigationOptionState::RtlMitigationOptionStateForce as u64;
        }

        if value == OptionValue::On && option == OptionValue::On {
            option_state |= RtlImageMitigationOptionState::RtlMitigationOptionStateOption as u64;
        }

        option_state
    }

    pub fn write_policy_to_registry(app_mits: &AppMitigations) -> Result<(), String> {
        let image_path = if app_mits.process_name.eq_ignore_ascii_case("System") {
            None
        } else {
            Some(app_mits.process_name.as_str())
        };

        // DEP Policy
        if let Some(AppPolicyContainer::Dep(dep_policy)) = app_mits.policies.get(&PolicyName::DEP) {
            let dep_struct = RtlImageMitigationDepPolicy {
                dep: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        dep_policy.get_policy(PolicyOptionName::Enable),
                        dep_policy.override_dep(),
                        dep_policy.get_policy(PolicyOptionName::EmulateAtlThunks),
                    ),
                },
            };

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageDepPolicy,
                dep_policy.flags(),
                &dep_struct,
            );

            if result != 0 {
                println!("Error while setting DEP: {:X}", result);
            }
        }

        // ASLR Policy
        if let Some(AppPolicyContainer::Aslr(aslr_policy)) =
            app_mits.policies.get(&PolicyName::ASLR)
        {
            let aslr_struct = RtlImageMitigationAslrPolicy {
                force_relocate_images: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        aslr_policy.get_policy(PolicyOptionName::ForceRelocateImages),
                        aslr_policy.override_force_relocate_images(),
                        aslr_policy.get_policy(PolicyOptionName::RequireInfo),
                    ),
                },
                bottom_up_randomization: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        aslr_policy.get_policy(PolicyOptionName::BottomUp),
                        aslr_policy.override_bottom_up(),
                        OptionValue::Off,
                    ),
                },
                high_entropy_randomization: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        aslr_policy.get_policy(PolicyOptionName::HighEntropy),
                        aslr_policy.override_high_entropy(),
                        OptionValue::Off,
                    ),
                },
            };

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageAslrPolicy,
                aslr_policy.flags(),
                &aslr_struct,
            );

            if result != 0 {
                println!("Error while setting ASLR: {:X}", result);
            }
        }

        // StrictHandle Policy
        if let Some(AppPolicyContainer::StrictHandle(strict_handle_policy)) =
            app_mits.policies.get(&PolicyName::StrictHandle)
        {
            let strict_handle_struct = RtlImageMitigationStrictHandleCheckPolicy {
                strict_handle_checks: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        strict_handle_policy.get_policy(PolicyOptionName::Enable),
                        strict_handle_policy.override_strict_handle(),
                        OptionValue::Off,
                    ),
                },
            };

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageStrictHandleCheckPolicy,
                strict_handle_policy.flags(),
                &strict_handle_struct,
            );

            if result != 0 {
                println!("Error while setting StrictHandle: {:X}", result);
            }
        }

        // SystemCall Policy
        if let Some(AppPolicyContainer::SystemCalls(system_call_policy)) =
            app_mits.policies.get(&PolicyName::SystemCalls)
        {
            let mut system_call_struct = RtlImageMitigationSystemCallDisablePolicy {
                block_win32k_system_calls: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        system_call_policy.get_policy(PolicyOptionName::DisableWin32kSystemCalls),
                        system_call_policy.override_system_call(),
                        OptionValue::Off,
                    ),
                },
                block_fsctl_system_calls: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        system_call_policy.get_policy(PolicyOptionName::DisableFsctlSystemCalls),
                        system_call_policy.override_fsctl_system_call(),
                        OptionValue::Off,
                    ),
                },
            };

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageSystemCallDisablePolicy,
                system_call_policy.flags(),
                &system_call_struct,
            );

            if result != 0 {
                println!("Error while setting SystemCall: {:X}", result);
            }

            // Audit flags
            system_call_struct.block_win32k_system_calls.policy_state =
                Self::option_value_to_option_state(
                    system_call_policy.get_policy(PolicyOptionName::Audit),
                    false,
                    OptionValue::Off,
                );
            system_call_struct.block_fsctl_system_calls.policy_state =
                Self::option_value_to_option_state(
                    system_call_policy.get_policy(PolicyOptionName::AuditFsctlSystemCalls),
                    false,
                    OptionValue::Off,
                );

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageSystemCallDisablePolicy,
                8 | system_call_policy.flags(),
                &system_call_struct,
            );

            if result != 0 {
                println!("Error while setting SystemCallAudit: {:X}", result);
            }
        }

        // ExtensionPoint Policy
        if let Some(AppPolicyContainer::ExtensionPoints(extension_point_policy)) =
            app_mits.policies.get(&PolicyName::ExtensionPoints)
        {
            let extension_point_struct = RtlImageMitigationExtensionPointDisablePolicy {
                disable_extension_points: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        extension_point_policy.get_policy(PolicyOptionName::DisableExtensionPoints),
                        extension_point_policy.override_extension_point(),
                        OptionValue::Off,
                    ),
                },
            };

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageExtensionPointDisablePolicy,
                extension_point_policy.flags(),
                &extension_point_struct,
            );

            if result != 0 {
                println!("Error while setting ExtensionPoint: {:X}", result);
            }
        }

        // DynamicCode Policy
        if let Some(AppPolicyContainer::DynamicCode(dynamic_code_policy)) =
            app_mits.policies.get(&PolicyName::DynamicCode)
        {
            let mut dynamic_code_struct = RtlImageMitigationDynamicCodePolicy {
                block_dynamic_code: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        dynamic_code_policy.get_policy(PolicyOptionName::BlockDynamicCode),
                        dynamic_code_policy.override_dynamic_code(),
                        OptionValue::Off,
                    ),
                },
            };

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageDynamicCodePolicy,
                dynamic_code_policy.flags(),
                &dynamic_code_struct,
            );

            if result != 0 {
                println!("Error while setting DynamicCode: {:X}", result);
            }

            // Audit flags
            dynamic_code_struct.block_dynamic_code.policy_state =
                Self::option_value_to_option_state(
                    dynamic_code_policy.get_policy(PolicyOptionName::Audit),
                    false,
                    OptionValue::Off,
                );

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageDynamicCodePolicy,
                8 | dynamic_code_policy.flags(),
                &dynamic_code_struct,
            );

            if result != 0 {
                println!("Error while setting DynamicCodeAudit: {:X}", result);
            }
        }

        // CFG Policy
        if let Some(AppPolicyContainer::ControlFlowGuard(cfg_policy)) =
            app_mits.policies.get(&PolicyName::ControlFlowGuard)
        {
            let cfg_struct = RtlImageMitigationControlFlowGuardPolicy {
                control_flow_guard: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        cfg_policy.get_policy(PolicyOptionName::Enable),
                        cfg_policy.override_cfg(),
                        cfg_policy.get_policy(PolicyOptionName::SuppressExports),
                    ),
                },
                strict_control_flow_guard: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        cfg_policy.get_policy(PolicyOptionName::StrictControlFlowGuard),
                        cfg_policy.override_strict_cfg(),
                        OptionValue::Off,
                    ),
                },
            };

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageControlFlowGuardPolicy,
                cfg_policy.flags(),
                &cfg_struct,
            );

            if result != 0 {
                println!("Error while setting CFG: {:X}", result);
            }
        }

        // BinarySignature Policy
        if let Some(AppPolicyContainer::SignedBinaries(binary_signature_policy)) =
            app_mits.policies.get(&PolicyName::SignedBinaries)
        {
            let mut binary_signature_struct = RtlImageMitigationBinarySignaturePolicy {
                block_non_microsoft_signed_binaries: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        binary_signature_policy.get_policy(PolicyOptionName::MicrosoftSignedOnly),
                        binary_signature_policy.override_microsoft_signed_only(),
                        binary_signature_policy
                            .get_policy(PolicyOptionName::AllowStoreSignedBinaries),
                    ),
                },
                enforce_signing_on_module_dependencies: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        binary_signature_policy
                            .get_policy(PolicyOptionName::EnforceModuleDependencySigning),
                        binary_signature_policy.override_enforce_module_dependency_signing(),
                        OptionValue::Off,
                    ),
                },
            };

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageSignaturePolicy,
                binary_signature_policy.flags(),
                &binary_signature_struct,
            );

            if result != 0 {
                println!("Error while setting BinarySignature: {:X}", result);
            }

            // Audit flags
            binary_signature_struct
                .block_non_microsoft_signed_binaries
                .policy_state = Self::option_value_to_option_state(
                binary_signature_policy.get_policy(PolicyOptionName::Audit),
                false,
                binary_signature_policy.get_policy(PolicyOptionName::AuditStoreSigned),
            );
            binary_signature_struct
                .enforce_signing_on_module_dependencies
                .policy_state = Self::option_value_to_option_state(
                binary_signature_policy
                    .get_policy(PolicyOptionName::AuditEnforceModuleDependencySigning),
                false,
                OptionValue::Off,
            );

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageSignaturePolicy,
                8 | binary_signature_policy.flags(),
                &binary_signature_struct,
            );

            if result != 0 {
                println!("Error while setting BinarySignatureAudit: {:X}", result);
            }
        }

        // FontDisable Policy
        if let Some(AppPolicyContainer::Fonts(font_disable_policy)) =
            app_mits.policies.get(&PolicyName::Fonts)
        {
            let mut font_disable_struct = RtlImageMitigationFontDisablePolicy {
                disable_non_system_fonts: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        font_disable_policy.get_policy(PolicyOptionName::DisableNonSystemFonts),
                        font_disable_policy.override_font_disable(),
                        OptionValue::Off,
                    ),
                },
            };

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageFontDisablePolicy,
                font_disable_policy.flags(),
                &font_disable_struct,
            );

            if result != 0 {
                println!("Error while setting FontDisable: {:X}", result);
            }

            // Audit flags
            font_disable_struct.disable_non_system_fonts.policy_state =
                Self::option_value_to_option_state(
                    font_disable_policy.get_policy(PolicyOptionName::Audit),
                    false,
                    OptionValue::Off,
                );

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageFontDisablePolicy,
                8 | font_disable_policy.flags(),
                &font_disable_struct,
            );

            if result != 0 {
                println!("Error while setting FontDisableAudit: {:X}", result);
            }
        }

        // ImageLoad Policy
        if let Some(AppPolicyContainer::ImageLoad(image_load_policy)) =
            app_mits.policies.get(&PolicyName::ImageLoad)
        {
            let mut image_load_struct = RtlImageMitigationImageLoadPolicy {
                block_remote_image_loads: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        image_load_policy.get_policy(PolicyOptionName::BlockRemoteImageLoads),
                        image_load_policy.override_block_remote_image_loads(),
                        OptionValue::Off,
                    ),
                },
                block_low_label_image_loads: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        image_load_policy.get_policy(PolicyOptionName::BlockLowLabelImageLoads),
                        image_load_policy.override_block_low_label(),
                        OptionValue::Off,
                    ),
                },
                prefer_system32: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        image_load_policy.get_policy(PolicyOptionName::PreferSystem32),
                        image_load_policy.override_prefer_system32(),
                        OptionValue::Off,
                    ),
                },
            };

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageImageLoadPolicy,
                image_load_policy.flags(),
                &image_load_struct,
            );

            if result != 0 {
                println!("Error while setting ImageLoad: {:X}", result);
            }

            // Audit flags
            image_load_struct.block_remote_image_loads.policy_state =
                Self::option_value_to_option_state(
                    image_load_policy.get_policy(PolicyOptionName::AuditRemoteImageLoads),
                    false,
                    OptionValue::Off,
                );
            image_load_struct.block_low_label_image_loads.policy_state =
                Self::option_value_to_option_state(
                    image_load_policy.get_policy(PolicyOptionName::AuditLowLabelImageLoads),
                    false,
                    OptionValue::Off,
                );
            image_load_struct.prefer_system32.policy_state = Self::option_value_to_option_state(
                image_load_policy.get_policy(PolicyOptionName::AuditPreferSystem32),
                false,
                OptionValue::Off,
            );

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageImageLoadPolicy,
                8 | image_load_policy.flags(),
                &image_load_struct,
            );

            if result != 0 {
                println!("Error while setting ImageLoadAudit: {:X}", result);
            }
        }

        // Payload Policy
        if let Some(AppPolicyContainer::Payload(payload_policy)) =
            app_mits.policies.get(&PolicyName::Payload)
        {
            let mut module_list = [0u16; 512];
            let mut flag = false;

            if !payload_policy.eaf_modules.is_empty() {
                if payload_policy.eaf_modules.len() > 32 {
                    println!(
                        "Error while setting Payload: Number of EAFPlus modules ({}) exceeded maximum length ({})",
                        payload_policy.eaf_modules.len(),
                        32
                    );
                    flag = true;
                }

                let mut str_builder = String::new();
                for module in &payload_policy.eaf_modules {
                    if !module.is_empty() {
                        str_builder.push_str(module);
                        str_builder.push(';');
                    }
                }

                if !str_builder.is_empty() {
                    let char_array: Vec<u16> = str_builder.encode_utf16().collect();
                    if char_array.len() > module_list.len() {
                        println!(
                            "Error while setting Payload: EAFPlus modules list length {} exceeded maximum length {}",
                            char_array.len(),
                            module_list.len()
                        );
                        flag = true;
                    } else {
                        module_list[..char_array.len()].copy_from_slice(&char_array);
                    }
                }
            }

            if !flag {
                let payload_struct = RtlImageMitigationPayloadRestrictionPolicy {
                    enable_export_address_filter: RtlImageMitigationPolicy {
                        policy_state: Self::option_value_to_option_state(
                            payload_policy.get_policy(PolicyOptionName::EnableExportAddressFilter),
                            payload_policy.override_enable_export_address_filter(),
                            OptionValue::Off,
                        ),
                    },
                    enable_export_address_filter_plus: RtlImageMitigationPolicy {
                        policy_state: Self::option_value_to_option_state(
                            payload_policy
                                .get_policy(PolicyOptionName::EnableExportAddressFilterPlus),
                            payload_policy.override_enable_export_address_filter_plus(),
                            OptionValue::Off,
                        ),
                    },
                    enable_import_address_filter: RtlImageMitigationPolicy {
                        policy_state: Self::option_value_to_option_state(
                            payload_policy.get_policy(PolicyOptionName::EnableImportAddressFilter),
                            payload_policy.override_enable_import_address_filter(),
                            OptionValue::Off,
                        ),
                    },
                    enable_rop_stack_pivot: RtlImageMitigationPolicy {
                        policy_state: Self::option_value_to_option_state(
                            payload_policy.get_policy(PolicyOptionName::EnableRopStackPivot),
                            payload_policy.override_enable_rop_stack_pivot(),
                            OptionValue::Off,
                        ),
                    },
                    enable_rop_caller_check: RtlImageMitigationPolicy {
                        policy_state: Self::option_value_to_option_state(
                            payload_policy.get_policy(PolicyOptionName::EnableRopCallerCheck),
                            payload_policy.override_enable_rop_caller_check(),
                            OptionValue::Off,
                        ),
                    },
                    enable_rop_sim_exec: RtlImageMitigationPolicy {
                        policy_state: Self::option_value_to_option_state(
                            payload_policy.get_policy(PolicyOptionName::EnableRopSimExec),
                            payload_policy.override_enable_rop_sim_exec(),
                            OptionValue::Off,
                        ),
                    },
                    module_list,
                };

                let result = Ntdll::set_image_mitigation_policy(
                    image_path,
                    ImageMitigationPolicy::ImagePayloadRestrictionPolicy,
                    payload_policy.flags(),
                    &payload_struct,
                );

                if result != 0 {
                    println!("Error while setting Payload: {:X}", result);
                }
            }

            // Audit flags
            let audit_payload_struct = RtlImageMitigationPayloadRestrictionPolicy {
                enable_export_address_filter: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        payload_policy.get_policy(PolicyOptionName::AuditEnableExportAddressFilter),
                        false,
                        OptionValue::Off,
                    ),
                },
                enable_export_address_filter_plus: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        payload_policy
                            .get_policy(PolicyOptionName::AuditEnableExportAddressFilterPlus),
                        false,
                        OptionValue::Off,
                    ),
                },
                enable_import_address_filter: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        payload_policy.get_policy(PolicyOptionName::AuditEnableImportAddressFilter),
                        false,
                        OptionValue::Off,
                    ),
                },
                enable_rop_stack_pivot: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        payload_policy.get_policy(PolicyOptionName::AuditEnableRopStackPivot),
                        false,
                        OptionValue::Off,
                    ),
                },
                enable_rop_caller_check: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        payload_policy.get_policy(PolicyOptionName::AuditEnableRopCallerCheck),
                        false,
                        OptionValue::Off,
                    ),
                },
                enable_rop_sim_exec: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        payload_policy.get_policy(PolicyOptionName::AuditEnableRopSimExec),
                        false,
                        OptionValue::Off,
                    ),
                },
                module_list: [0; 512],
            };

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImagePayloadRestrictionPolicy,
                8 | payload_policy.flags(),
                &audit_payload_struct,
            );

            if result != 0 {
                println!("Error while setting PayloadAudit: {:X}", result);
            }
        }

        // ChildProcess Policy
        if let Some(AppPolicyContainer::ChildProcess(child_process_policy)) =
            app_mits.policies.get(&PolicyName::ChildProcess)
        {
            let mut child_process_struct = RtlImageMitigationChildProcessPolicy {
                disallow_child_process_creation: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        child_process_policy
                            .get_policy(PolicyOptionName::DisallowChildProcessCreation),
                        child_process_policy.override_child_process(),
                        OptionValue::Off,
                    ),
                },
            };

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageChildProcessPolicy,
                child_process_policy.flags(),
                &child_process_struct,
            );

            if result != 0 {
                println!("Error while setting ChildProcess: {:X}", result);
            }

            // Audit flags
            child_process_struct
                .disallow_child_process_creation
                .policy_state = Self::option_value_to_option_state(
                child_process_policy.get_policy(PolicyOptionName::Audit),
                false,
                OptionValue::Off,
            );

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageChildProcessPolicy,
                8 | child_process_policy.flags(),
                &child_process_struct,
            );

            if result != 0 {
                println!("Error while setting ChildProcessAudit: {:X}", result);
            }
        }

        // SEHOP Policy
        if let Some(AppPolicyContainer::SEHOP(sehop_policy)) =
            app_mits.policies.get(&PolicyName::SEHOP)
        {
            let mut sehop_struct = RtlImageMitigationSehopPolicy {
                sehop: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        sehop_policy.get_policy(PolicyOptionName::Enable),
                        sehop_policy.override_sehop(),
                        sehop_policy.get_policy(PolicyOptionName::TelemetryOnly),
                    ),
                },
            };

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageSehopPolicy,
                sehop_policy.flags(),
                &sehop_struct,
            );

            if result != 0 {
                println!("Error while setting SEHOP: {:X}", result);
            }

            // Audit flags
            sehop_struct.sehop.policy_state = Self::option_value_to_option_state(
                sehop_policy.get_policy(PolicyOptionName::Audit),
                false,
                OptionValue::Off,
            );

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageSehopPolicy,
                8 | sehop_policy.flags(),
                &sehop_struct,
            );

            if result != 0 {
                println!("Error while setting SEHOPAudit: {:X}", result);
            }
        }

        // Heap Policy
        if let Some(AppPolicyContainer::Heap(heap_policy)) =
            app_mits.policies.get(&PolicyName::Heap)
        {
            let heap_struct = RtlImageMitigationHeapPolicy {
                terminate_on_heap_errors: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        heap_policy.get_policy(PolicyOptionName::TerminateOnError),
                        heap_policy.override_heap(),
                        OptionValue::Off,
                    ),
                },
            };

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageHeapPolicy,
                heap_policy.flags(),
                &heap_struct,
            );

            if result != 0 {
                println!("Error while setting Heap: {:X}", result);
            }
        }

        // UserShadowStack Policy
        if let Some(AppPolicyContainer::UserShadowStack(user_shadow_stack_policy)) =
            app_mits.policies.get(&PolicyName::UserShadowStack)
        {
            let mut user_shadow_stack_struct = RtlImageMitigationUserShadowStackPolicy {
                user_shadow_stack: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        user_shadow_stack_policy.get_policy(PolicyOptionName::UserShadowStack),
                        user_shadow_stack_policy.override_user_shadow_stack(),
                        user_shadow_stack_policy
                            .get_policy(PolicyOptionName::UserShadowStackStrictMode),
                    ),
                },
                set_context_ip_validation: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        user_shadow_stack_policy
                            .get_policy(PolicyOptionName::SetContextIpValidation),
                        false,
                        OptionValue::Off,
                    ),
                },
                block_non_cet_binaries: RtlImageMitigationPolicy {
                    policy_state: Self::option_value_to_option_state(
                        user_shadow_stack_policy.get_policy(PolicyOptionName::BlockNonCetBinaries),
                        false,
                        OptionValue::Off,
                    ),
                },
            };

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageUserShadowStackPolicy,
                user_shadow_stack_policy.flags(),
                &user_shadow_stack_struct,
            );

            if result != 0 {
                println!("Error while setting UserShadowStack: {:X}", result);
            }

            // Audit flags
            user_shadow_stack_struct.user_shadow_stack.policy_state =
                Self::option_value_to_option_state(
                    user_shadow_stack_policy.get_policy(PolicyOptionName::AuditUserShadowStack),
                    false,
                    OptionValue::Off,
                );

            let result = Ntdll::set_image_mitigation_policy(
                image_path,
                ImageMitigationPolicy::ImageUserShadowStackPolicy,
                8 | user_shadow_stack_policy.flags(),
                &user_shadow_stack_struct,
            );

            if result != 0 {
                println!("Error while setting UserShadowStackAudit: {:X}", result);
            }
        }

        Ok(())
    }

    fn option_state_to_option_value(
        state: u64,
        force: &mut dyn FnMut(bool),
        option: &mut dyn FnMut(OptionValue),
    ) -> OptionValue {
        let option_value = match state & 3 {
            1 => OptionValue::On,
            2 => OptionValue::Off,
            _ => OptionValue::NotSet,
        };
        force((state & 4) > 0);
        option(if (state & 8) != 0 {
            OptionValue::On
        } else {
            OptionValue::Off
        });
        option_value
    }

    fn build_full_policy(exe_name: Option<&str>) -> Result<AppMitigations, String> {
        let mut mits = AppMitigations::new(exe_name.unwrap_or(""));

        // Query DEP Policy
        match Ntdll::query_image_mitigation_policy::<RtlImageMitigationDepPolicy>(
            exe_name,
            ImageMitigationPolicy::ImageDepPolicy,
            0,
        ) {
            Ok(policy) => {
                if let Some(AppPolicyContainer::Dep(dep_policy)) =
                    mits.policies.get_mut(&PolicyName::DEP)
                {
                    let mut force_value = false;
                    let mut option_value = OptionValue::Off;

                    let enable_value = Self::option_state_to_option_value(
                        policy.dep.policy_state,
                        &mut |force| force_value = force,
                        &mut |option| option_value = option,
                    );

                    dep_policy.set_override_dep(force_value);
                    dep_policy.set_policy(PolicyOptionName::EmulateAtlThunks, option_value);
                    dep_policy.set_policy(PolicyOptionName::Enable, enable_value);
                }
            }
            Err(-1073741772) => {
                // STATUS_OBJECT_NAME_NOT_FOUND - continue to next policy
            }
            Err(code) => {
                println!(
                    "Warning: there was an error querying DEP policy: {:X}",
                    code
                );
            }
        }

        // Query ASLR Policy
        match Ntdll::query_image_mitigation_policy::<RtlImageMitigationAslrPolicy>(
            exe_name,
            ImageMitigationPolicy::ImageAslrPolicy,
            0,
        ) {
            Ok(policy) => {
                if let Some(AppPolicyContainer::Aslr(aslr_policy)) =
                    mits.policies.get_mut(&PolicyName::ASLR)
                {
                    let mut force_relocate_force = false;
                    let mut require_info_option = OptionValue::Off;
                    let force_relocate_value = Self::option_state_to_option_value(
                        policy.force_relocate_images.policy_state,
                        &mut |force| force_relocate_force = force,
                        &mut |option| require_info_option = option,
                    );
                    aslr_policy.set_override_force_relocate_images(force_relocate_force);
                    aslr_policy.set_policy(PolicyOptionName::RequireInfo, require_info_option);
                    aslr_policy
                        .set_policy(PolicyOptionName::ForceRelocateImages, force_relocate_value);

                    let mut bottom_up_force = false;
                    let bottom_up_value = Self::option_state_to_option_value(
                        policy.bottom_up_randomization.policy_state,
                        &mut |force| bottom_up_force = force,
                        &mut |_| {},
                    );
                    aslr_policy.set_override_bottom_up(bottom_up_force);
                    aslr_policy.set_policy(PolicyOptionName::BottomUp, bottom_up_value);

                    let mut high_entropy_force = false;
                    let high_entropy_value = Self::option_state_to_option_value(
                        policy.high_entropy_randomization.policy_state,
                        &mut |force| high_entropy_force = force,
                        &mut |_| {},
                    );
                    aslr_policy.set_override_high_entropy(high_entropy_force);
                    aslr_policy.set_policy(PolicyOptionName::HighEntropy, high_entropy_value);
                }
            }
            Err(-1073741772) => {
                // STATUS_OBJECT_NAME_NOT_FOUND - continue to next policy
            }
            Err(code) => {
                println!("Warning: error while querying ASLR policy: {:X}", code);
            }
        }

        // Query StrictHandle Policy
        match Ntdll::query_image_mitigation_policy::<RtlImageMitigationStrictHandleCheckPolicy>(
            exe_name,
            ImageMitigationPolicy::ImageStrictHandleCheckPolicy,
            0,
        ) {
            Ok(policy) => {
                if let Some(AppPolicyContainer::StrictHandle(strict_handle_policy)) =
                    mits.policies.get_mut(&PolicyName::StrictHandle)
                {
                    let mut force_value = false;
                    let enable_value = Self::option_state_to_option_value(
                        policy.strict_handle_checks.policy_state,
                        &mut |force| force_value = force,
                        &mut |_| {},
                    );
                    strict_handle_policy.set_override_strict_handle(force_value);
                    strict_handle_policy.set_policy(PolicyOptionName::Enable, enable_value);
                }
            }
            Err(-1073741772) => {
                // STATUS_OBJECT_NAME_NOT_FOUND - continue to next policy
            }
            Err(code) => {
                println!(
                    "Warning: error while querying StrictHandle policy: {:X}",
                    code
                );
            }
        }

        // Query SystemCall Policy
        match Ntdll::query_image_mitigation_policy::<RtlImageMitigationSystemCallDisablePolicy>(
            exe_name,
            ImageMitigationPolicy::ImageSystemCallDisablePolicy,
            0,
        ) {
            Ok(policy) => {
                if let Some(AppPolicyContainer::SystemCalls(system_call_policy)) =
                    mits.policies.get_mut(&PolicyName::SystemCalls)
                {
                    let mut win32k_force = false;
                    let win32k_value = Self::option_state_to_option_value(
                        policy.block_win32k_system_calls.policy_state,
                        &mut |force| win32k_force = force,
                        &mut |_| {},
                    );
                    system_call_policy.set_override_system_call(win32k_force);
                    system_call_policy
                        .set_policy(PolicyOptionName::DisableWin32kSystemCalls, win32k_value);

                    let mut fsctl_force = false;
                    let fsctl_value = Self::option_state_to_option_value(
                        policy.block_fsctl_system_calls.policy_state,
                        &mut |force| fsctl_force = force,
                        &mut |_| {},
                    );
                    system_call_policy.set_override_fsctl_system_call(fsctl_force);
                    system_call_policy
                        .set_policy(PolicyOptionName::DisableFsctlSystemCalls, fsctl_value);
                }

                // Query audit flags
                match Ntdll::query_image_mitigation_policy::<
                    RtlImageMitigationSystemCallDisablePolicy,
                >(
                    exe_name,
                    ImageMitigationPolicy::ImageSystemCallDisablePolicy,
                    8,
                ) {
                    Ok(audit_policy) => {
                        if let Some(AppPolicyContainer::SystemCalls(system_call_policy)) =
                            mits.policies.get_mut(&PolicyName::SystemCalls)
                        {
                            let audit_value = Self::option_state_to_option_value(
                                audit_policy.block_win32k_system_calls.policy_state,
                                &mut |_| {},
                                &mut |_| {},
                            );
                            system_call_policy.set_policy(PolicyOptionName::Audit, audit_value);

                            let audit_fsctl_value = Self::option_state_to_option_value(
                                audit_policy.block_fsctl_system_calls.policy_state,
                                &mut |_| {},
                                &mut |_| {},
                            );
                            system_call_policy.set_policy(
                                PolicyOptionName::AuditFsctlSystemCalls,
                                audit_fsctl_value,
                            );
                        }
                    }
                    Err(code) => {
                        println!(
                            "Warning: error while querying Syscall audit policy: {:X}",
                            code
                        );
                    }
                }
            }
            Err(-1073741772) => {
                // STATUS_OBJECT_NAME_NOT_FOUND - continue to next policy
            }
            Err(code) => {
                println!("Warning: error while querying Syscall policy: {:X}", code);
            }
        }

        // Query ExtensionPoint Policy
        match Ntdll::query_image_mitigation_policy::<RtlImageMitigationExtensionPointDisablePolicy>(
            exe_name,
            ImageMitigationPolicy::ImageExtensionPointDisablePolicy,
            0,
        ) {
            Ok(policy) => {
                if let Some(AppPolicyContainer::ExtensionPoints(extension_point_policy)) =
                    mits.policies.get_mut(&PolicyName::ExtensionPoints)
                {
                    let mut force_value = false;
                    let value = Self::option_state_to_option_value(
                        policy.disable_extension_points.policy_state,
                        &mut |force| force_value = force,
                        &mut |_| {},
                    );
                    extension_point_policy.set_override_extension_point(force_value);
                    extension_point_policy
                        .set_policy(PolicyOptionName::DisableExtensionPoints, value);
                }
            }
            Err(-1073741772) => {}
            Err(code) => {
                println!(
                    "Warning: error while querying ExtensionPoint policy: {:X}",
                    code
                );
            }
        }

        // Query DynamicCode Policy
        match Ntdll::query_image_mitigation_policy::<RtlImageMitigationDynamicCodePolicy>(
            exe_name,
            ImageMitigationPolicy::ImageDynamicCodePolicy,
            0,
        ) {
            Ok(policy) => {
                if let Some(AppPolicyContainer::DynamicCode(dynamic_code_policy)) =
                    mits.policies.get_mut(&PolicyName::DynamicCode)
                {
                    let mut force_value = false;
                    let value = Self::option_state_to_option_value(
                        policy.block_dynamic_code.policy_state,
                        &mut |force| force_value = force,
                        &mut |_| {},
                    );
                    dynamic_code_policy.set_override_dynamic_code(force_value);
                    dynamic_code_policy.set_policy(PolicyOptionName::BlockDynamicCode, value);
                    // AllowThreadsToOptOut remains NotSet as it's not set from BlockDynamicCode policy
                }

                // Query audit flags
                match Ntdll::query_image_mitigation_policy::<RtlImageMitigationDynamicCodePolicy>(
                    exe_name,
                    ImageMitigationPolicy::ImageDynamicCodePolicy,
                    8,
                ) {
                    Ok(audit_policy) => {
                        if let Some(AppPolicyContainer::DynamicCode(dynamic_code_policy)) =
                            mits.policies.get_mut(&PolicyName::DynamicCode)
                        {
                            let audit_value = Self::option_state_to_option_value(
                                audit_policy.block_dynamic_code.policy_state,
                                &mut |_| {},
                                &mut |_| {},
                            );
                            dynamic_code_policy.set_policy(PolicyOptionName::Audit, audit_value);
                        }
                    }
                    Err(code) => {
                        println!(
                            "Warning: error while querying DynamicCode Audit policy: {:X}",
                            code
                        );
                    }
                }
            }
            Err(-1073741772) => {}
            Err(code) => {
                println!(
                    "Warning: error while querying DynamicCode policy: {:X}",
                    code
                );
            }
        }

        // Query CFG Policy
        match Ntdll::query_image_mitigation_policy::<RtlImageMitigationControlFlowGuardPolicy>(
            exe_name,
            ImageMitigationPolicy::ImageControlFlowGuardPolicy,
            0,
        ) {
            Ok(policy) => {
                if let Some(AppPolicyContainer::ControlFlowGuard(cfg_policy)) =
                    mits.policies.get_mut(&PolicyName::ControlFlowGuard)
                {
                    let mut cfg_force = false;
                    let mut suppress_exports_option = OptionValue::Off;
                    let enable_value = Self::option_state_to_option_value(
                        policy.control_flow_guard.policy_state,
                        &mut |force| cfg_force = force,
                        &mut |option| suppress_exports_option = option,
                    );
                    cfg_policy.set_override_cfg(cfg_force);
                    cfg_policy
                        .set_policy(PolicyOptionName::SuppressExports, suppress_exports_option);
                    cfg_policy.set_policy(PolicyOptionName::Enable, enable_value);

                    let mut strict_force = false;
                    let strict_value = Self::option_state_to_option_value(
                        policy.strict_control_flow_guard.policy_state,
                        &mut |force| strict_force = force,
                        &mut |_| {},
                    );
                    cfg_policy.set_override_strict_cfg(strict_force);
                    cfg_policy.set_policy(PolicyOptionName::StrictControlFlowGuard, strict_value);
                }
            }
            Err(-1073741772) => {}
            Err(code) => {
                println!("Warning: error while querying CFG policy: {:X}", code);
            }
        }

        // Query BinarySignature Policy
        match Ntdll::query_image_mitigation_policy::<RtlImageMitigationBinarySignaturePolicy>(
            exe_name,
            ImageMitigationPolicy::ImageSignaturePolicy,
            0,
        ) {
            Ok(policy) => {
                if let Some(AppPolicyContainer::SignedBinaries(binary_signature_policy)) =
                    mits.policies.get_mut(&PolicyName::SignedBinaries)
                {
                    let mut ms_force = false;
                    let mut store_signed_option = OptionValue::Off;
                    let ms_value = Self::option_state_to_option_value(
                        policy.block_non_microsoft_signed_binaries.policy_state,
                        &mut |force| ms_force = force,
                        &mut |option| store_signed_option = option,
                    );
                    binary_signature_policy.set_override_microsoft_signed_only(ms_force);
                    binary_signature_policy.set_policy(
                        PolicyOptionName::AllowStoreSignedBinaries,
                        store_signed_option,
                    );
                    binary_signature_policy
                        .set_policy(PolicyOptionName::MicrosoftSignedOnly, ms_value);

                    let mut enforce_force = false;
                    let enforce_value = Self::option_state_to_option_value(
                        policy.enforce_signing_on_module_dependencies.policy_state,
                        &mut |force| enforce_force = force,
                        &mut |_| {},
                    );
                    binary_signature_policy
                        .set_override_enforce_module_dependency_signing(enforce_force);
                    binary_signature_policy.set_policy(
                        PolicyOptionName::EnforceModuleDependencySigning,
                        enforce_value,
                    );
                }

                // Query audit flags
                match Ntdll::query_image_mitigation_policy::<RtlImageMitigationBinarySignaturePolicy>(
                    exe_name,
                    ImageMitigationPolicy::ImageSignaturePolicy,
                    8,
                ) {
                    Ok(audit_policy) => {
                        if let Some(AppPolicyContainer::SignedBinaries(binary_signature_policy)) =
                            mits.policies.get_mut(&PolicyName::SignedBinaries)
                        {
                            let mut audit_store_option = OptionValue::Off;
                            let audit_value = Self::option_state_to_option_value(
                                audit_policy
                                    .block_non_microsoft_signed_binaries
                                    .policy_state,
                                &mut |_| {},
                                &mut |option| audit_store_option = option,
                            );
                            binary_signature_policy
                                .set_policy(PolicyOptionName::AuditStoreSigned, audit_store_option);
                            binary_signature_policy
                                .set_policy(PolicyOptionName::Audit, audit_value);

                            let audit_enforce_value = Self::option_state_to_option_value(
                                audit_policy
                                    .enforce_signing_on_module_dependencies
                                    .policy_state,
                                &mut |_| {},
                                &mut |_| {},
                            );
                            binary_signature_policy.set_policy(
                                PolicyOptionName::AuditEnforceModuleDependencySigning,
                                audit_enforce_value,
                            );
                        }
                    }
                    Err(code) => {
                        println!(
                            "Warning: error while querying SignedBinaries Audit policy: {:X}",
                            code
                        );
                    }
                }
            }
            Err(-1073741772) => {}
            Err(code) => {
                println!(
                    "Warning: error while querying SignedBinaries policy: {:X}",
                    code
                );
            }
        }

        // Query FontDisable Policy
        match Ntdll::query_image_mitigation_policy::<RtlImageMitigationFontDisablePolicy>(
            exe_name,
            ImageMitigationPolicy::ImageFontDisablePolicy,
            0,
        ) {
            Ok(policy) => {
                if let Some(AppPolicyContainer::Fonts(font_disable_policy)) =
                    mits.policies.get_mut(&PolicyName::Fonts)
                {
                    let mut force_value = false;
                    let value = Self::option_state_to_option_value(
                        policy.disable_non_system_fonts.policy_state,
                        &mut |force| force_value = force,
                        &mut |_| {},
                    );
                    font_disable_policy.set_override_font_disable(force_value);
                    font_disable_policy.set_policy(PolicyOptionName::DisableNonSystemFonts, value);
                }

                // Query audit flags
                match Ntdll::query_image_mitigation_policy::<RtlImageMitigationFontDisablePolicy>(
                    exe_name,
                    ImageMitigationPolicy::ImageFontDisablePolicy,
                    8,
                ) {
                    Ok(audit_policy) => {
                        if let Some(AppPolicyContainer::Fonts(font_disable_policy)) =
                            mits.policies.get_mut(&PolicyName::Fonts)
                        {
                            let audit_value = Self::option_state_to_option_value(
                                audit_policy.disable_non_system_fonts.policy_state,
                                &mut |_| {},
                                &mut |_| {},
                            );
                            font_disable_policy.set_policy(PolicyOptionName::Audit, audit_value);
                        }
                    }
                    Err(code) => {
                        println!(
                            "Warning: error while querying FontDisable Audit policy: {:X}",
                            code
                        );
                    }
                }
            }
            Err(-1073741772) => {}
            Err(code) => {
                println!(
                    "Warning: error while querying FontDisable policy: {:X}",
                    code
                );
            }
        }

        // Query ImageLoad Policy
        match Ntdll::query_image_mitigation_policy::<RtlImageMitigationImageLoadPolicy>(
            exe_name,
            ImageMitigationPolicy::ImageImageLoadPolicy,
            0,
        ) {
            Ok(policy) => {
                if let Some(AppPolicyContainer::ImageLoad(image_load_policy)) =
                    mits.policies.get_mut(&PolicyName::ImageLoad)
                {
                    let mut remote_force = false;
                    let remote_value = Self::option_state_to_option_value(
                        policy.block_remote_image_loads.policy_state,
                        &mut |force| remote_force = force,
                        &mut |_| {},
                    );
                    image_load_policy.set_override_block_remote_image_loads(remote_force);
                    image_load_policy
                        .set_policy(PolicyOptionName::BlockRemoteImageLoads, remote_value);

                    let mut low_label_force = false;
                    let low_label_value = Self::option_state_to_option_value(
                        policy.block_low_label_image_loads.policy_state,
                        &mut |force| low_label_force = force,
                        &mut |_| {},
                    );
                    image_load_policy.set_override_block_low_label(low_label_force);
                    image_load_policy
                        .set_policy(PolicyOptionName::BlockLowLabelImageLoads, low_label_value);

                    let mut prefer_sys32_force = false;
                    let prefer_sys32_value = Self::option_state_to_option_value(
                        policy.prefer_system32.policy_state,
                        &mut |force| prefer_sys32_force = force,
                        &mut |_| {},
                    );
                    image_load_policy.set_override_prefer_system32(prefer_sys32_force);
                    image_load_policy
                        .set_policy(PolicyOptionName::PreferSystem32, prefer_sys32_value);
                }

                // Query audit flags
                match Ntdll::query_image_mitigation_policy::<RtlImageMitigationImageLoadPolicy>(
                    exe_name,
                    ImageMitigationPolicy::ImageImageLoadPolicy,
                    8,
                ) {
                    Ok(audit_policy) => {
                        if let Some(AppPolicyContainer::ImageLoad(image_load_policy)) =
                            mits.policies.get_mut(&PolicyName::ImageLoad)
                        {
                            let audit_remote_value = Self::option_state_to_option_value(
                                audit_policy.block_remote_image_loads.policy_state,
                                &mut |_| {},
                                &mut |_| {},
                            );
                            image_load_policy.set_policy(
                                PolicyOptionName::AuditRemoteImageLoads,
                                audit_remote_value,
                            );

                            let audit_low_label_value = Self::option_state_to_option_value(
                                audit_policy.block_low_label_image_loads.policy_state,
                                &mut |_| {},
                                &mut |_| {},
                            );
                            image_load_policy.set_policy(
                                PolicyOptionName::AuditLowLabelImageLoads,
                                audit_low_label_value,
                            );

                            let audit_prefer_sys32_value = Self::option_state_to_option_value(
                                audit_policy.prefer_system32.policy_state,
                                &mut |_| {},
                                &mut |_| {},
                            );
                            image_load_policy.set_policy(
                                PolicyOptionName::AuditPreferSystem32,
                                audit_prefer_sys32_value,
                            );
                        }
                    }
                    Err(code) => {
                        println!(
                            "Warning: error while querying ImageLoad Audit policy: {:X}",
                            code
                        );
                    }
                }
            }
            Err(-1073741772) => {}
            Err(code) => {
                println!("Warning: error while querying ImageLoad policy: {:X}", code);
            }
        }

        // Query Payload Policy
        match Ntdll::query_image_mitigation_policy::<RtlImageMitigationPayloadRestrictionPolicy>(
            exe_name,
            ImageMitigationPolicy::ImagePayloadRestrictionPolicy,
            0,
        ) {
            Ok(policy) => {
                if let Some(AppPolicyContainer::Payload(payload_policy)) =
                    mits.policies.get_mut(&PolicyName::Payload)
                {
                    let mut eaf_force = false;
                    let eaf_value = Self::option_state_to_option_value(
                        policy.enable_export_address_filter.policy_state,
                        &mut |force| eaf_force = force,
                        &mut |_| {},
                    );
                    payload_policy.set_override_enable_export_address_filter(eaf_force);
                    payload_policy
                        .set_policy(PolicyOptionName::EnableExportAddressFilter, eaf_value);

                    let mut eaf_plus_force = false;
                    let eaf_plus_value = Self::option_state_to_option_value(
                        policy.enable_export_address_filter_plus.policy_state,
                        &mut |force| eaf_plus_force = force,
                        &mut |_| {},
                    );
                    payload_policy.set_override_enable_export_address_filter_plus(eaf_plus_force);
                    payload_policy.set_policy(
                        PolicyOptionName::EnableExportAddressFilterPlus,
                        eaf_plus_value,
                    );

                    let mut iaf_force = false;
                    let iaf_value = Self::option_state_to_option_value(
                        policy.enable_import_address_filter.policy_state,
                        &mut |force| iaf_force = force,
                        &mut |_| {},
                    );
                    payload_policy.set_override_enable_import_address_filter(iaf_force);
                    payload_policy
                        .set_policy(PolicyOptionName::EnableImportAddressFilter, iaf_value);

                    let mut rop_stack_force = false;
                    let rop_stack_value = Self::option_state_to_option_value(
                        policy.enable_rop_stack_pivot.policy_state,
                        &mut |force| rop_stack_force = force,
                        &mut |_| {},
                    );
                    payload_policy.set_override_enable_rop_stack_pivot(rop_stack_force);
                    payload_policy
                        .set_policy(PolicyOptionName::EnableRopStackPivot, rop_stack_value);

                    let mut rop_caller_force = false;
                    let rop_caller_value = Self::option_state_to_option_value(
                        policy.enable_rop_caller_check.policy_state,
                        &mut |force| rop_caller_force = force,
                        &mut |_| {},
                    );
                    payload_policy.set_override_enable_rop_caller_check(rop_caller_force);
                    payload_policy
                        .set_policy(PolicyOptionName::EnableRopCallerCheck, rop_caller_value);

                    let mut rop_sim_force = false;
                    let rop_sim_value = Self::option_state_to_option_value(
                        policy.enable_rop_sim_exec.policy_state,
                        &mut |force| rop_sim_force = force,
                        &mut |_| {},
                    );
                    payload_policy.set_override_enable_rop_sim_exec(rop_sim_force);
                    payload_policy.set_policy(PolicyOptionName::EnableRopSimExec, rop_sim_value);

                    // Parse module list
                    let str1 = String::from_utf16_lossy(&policy.module_list)
                        .trim_end_matches('\0')
                        .to_string();
                    if !str1.is_empty() {
                        let ch_array = [';'];
                        payload_policy.eaf_modules.clear();
                        for str2 in str1.split(&ch_array[..]) {
                            if !str2.is_empty()
                                && !payload_policy.eaf_modules.contains(&str2.to_string())
                            {
                                payload_policy.eaf_modules.push(str2.to_string());
                            }
                        }
                    }
                }

                // Query audit flags
                match Ntdll::query_image_mitigation_policy::<
                    RtlImageMitigationPayloadRestrictionPolicy,
                >(
                    exe_name,
                    ImageMitigationPolicy::ImagePayloadRestrictionPolicy,
                    8,
                ) {
                    Ok(audit_policy) => {
                        if let Some(AppPolicyContainer::Payload(payload_policy)) =
                            mits.policies.get_mut(&PolicyName::Payload)
                        {
                            let audit_eaf_value = Self::option_state_to_option_value(
                                audit_policy.enable_export_address_filter.policy_state,
                                &mut |_| {},
                                &mut |_| {},
                            );
                            payload_policy.set_policy(
                                PolicyOptionName::AuditEnableExportAddressFilter,
                                audit_eaf_value,
                            );

                            let audit_eaf_plus_value = Self::option_state_to_option_value(
                                audit_policy.enable_export_address_filter_plus.policy_state,
                                &mut |_| {},
                                &mut |_| {},
                            );
                            payload_policy.set_policy(
                                PolicyOptionName::AuditEnableExportAddressFilterPlus,
                                audit_eaf_plus_value,
                            );

                            let audit_iaf_value = Self::option_state_to_option_value(
                                audit_policy.enable_import_address_filter.policy_state,
                                &mut |_| {},
                                &mut |_| {},
                            );
                            payload_policy.set_policy(
                                PolicyOptionName::AuditEnableImportAddressFilter,
                                audit_iaf_value,
                            );

                            let audit_rop_stack_value = Self::option_state_to_option_value(
                                audit_policy.enable_rop_stack_pivot.policy_state,
                                &mut |_| {},
                                &mut |_| {},
                            );
                            payload_policy.set_policy(
                                PolicyOptionName::AuditEnableRopStackPivot,
                                audit_rop_stack_value,
                            );

                            let audit_rop_caller_value = Self::option_state_to_option_value(
                                audit_policy.enable_rop_caller_check.policy_state,
                                &mut |_| {},
                                &mut |_| {},
                            );
                            payload_policy.set_policy(
                                PolicyOptionName::AuditEnableRopCallerCheck,
                                audit_rop_caller_value,
                            );

                            let audit_rop_sim_value = Self::option_state_to_option_value(
                                audit_policy.enable_rop_sim_exec.policy_state,
                                &mut |_| {},
                                &mut |_| {},
                            );
                            payload_policy.set_policy(
                                PolicyOptionName::AuditEnableRopSimExec,
                                audit_rop_sim_value,
                            );
                        }
                    }
                    Err(code) => {
                        println!(
                            "Warning: error while querying Payload audit policy: {:X}",
                            code
                        );
                    }
                }
            }
            Err(-1073741772) => {}
            Err(code) => {
                println!("Warning: error while querying Payload policy: {:X}", code);
            }
        }

        // Query SEHOP Policy
        match Ntdll::query_image_mitigation_policy::<RtlImageMitigationSehopPolicy>(
            exe_name,
            ImageMitigationPolicy::ImageSehopPolicy,
            0,
        ) {
            Ok(policy) => {
                if let Some(AppPolicyContainer::SEHOP(sehop_policy)) =
                    mits.policies.get_mut(&PolicyName::SEHOP)
                {
                    let mut force_value = false;
                    let mut telemetry_option = OptionValue::Off;
                    let value = Self::option_state_to_option_value(
                        policy.sehop.policy_state,
                        &mut |force| force_value = force,
                        &mut |option| telemetry_option = option,
                    );
                    sehop_policy.set_override_sehop(force_value);
                    sehop_policy.set_policy(PolicyOptionName::TelemetryOnly, telemetry_option);
                    sehop_policy.set_policy(PolicyOptionName::Enable, value);
                }

                // Query audit flags
                match Ntdll::query_image_mitigation_policy::<RtlImageMitigationSehopPolicy>(
                    exe_name,
                    ImageMitigationPolicy::ImageSehopPolicy,
                    8,
                ) {
                    Ok(audit_policy) => {
                        if let Some(AppPolicyContainer::SEHOP(sehop_policy)) =
                            mits.policies.get_mut(&PolicyName::SEHOP)
                        {
                            let audit_value = Self::option_state_to_option_value(
                                audit_policy.sehop.policy_state,
                                &mut |_| {},
                                &mut |_| {},
                            );
                            sehop_policy.set_policy(PolicyOptionName::Audit, audit_value);
                        }
                    }
                    Err(code) => {
                        println!(
                            "Warning: error while querying SEHOP Audit policy: {:X}",
                            code
                        );
                    }
                }
            }
            Err(-1073741772) => {}
            Err(code) => {
                println!("Warning: error while querying SEHOP policy: {:X}", code);
            }
        }

        // Query Heap Policy
        match Ntdll::query_image_mitigation_policy::<RtlImageMitigationHeapPolicy>(
            exe_name,
            ImageMitigationPolicy::ImageHeapPolicy,
            0,
        ) {
            Ok(policy) => {
                if let Some(AppPolicyContainer::Heap(heap_policy)) =
                    mits.policies.get_mut(&PolicyName::Heap)
                {
                    let mut force_value = false;
                    let value = Self::option_state_to_option_value(
                        policy.terminate_on_heap_errors.policy_state,
                        &mut |force| force_value = force,
                        &mut |_| {},
                    );
                    heap_policy.set_override_heap(force_value);
                    heap_policy.set_policy(PolicyOptionName::TerminateOnError, value);
                }
            }
            Err(-1073741772) => {}
            Err(code) => {
                println!("Warning: error while querying Heap policy: {:X}", code);
            }
        }

        // Query ChildProcess Policy
        match Ntdll::query_image_mitigation_policy::<RtlImageMitigationChildProcessPolicy>(
            exe_name,
            ImageMitigationPolicy::ImageChildProcessPolicy,
            0,
        ) {
            Ok(policy) => {
                if let Some(AppPolicyContainer::ChildProcess(child_process_policy)) =
                    mits.policies.get_mut(&PolicyName::ChildProcess)
                {
                    let mut force_value = false;
                    let value = Self::option_state_to_option_value(
                        policy.disallow_child_process_creation.policy_state,
                        &mut |force| force_value = force,
                        &mut |_| {},
                    );
                    child_process_policy.set_override_child_process(force_value);
                    child_process_policy
                        .set_policy(PolicyOptionName::DisallowChildProcessCreation, value);
                }

                // Query audit flags
                match Ntdll::query_image_mitigation_policy::<RtlImageMitigationChildProcessPolicy>(
                    exe_name,
                    ImageMitigationPolicy::ImageChildProcessPolicy,
                    8,
                ) {
                    Ok(audit_policy) => {
                        if let Some(AppPolicyContainer::ChildProcess(child_process_policy)) =
                            mits.policies.get_mut(&PolicyName::ChildProcess)
                        {
                            let audit_value = Self::option_state_to_option_value(
                                audit_policy.disallow_child_process_creation.policy_state,
                                &mut |_| {},
                                &mut |_| {},
                            );
                            child_process_policy.set_policy(PolicyOptionName::Audit, audit_value);
                        }
                    }
                    Err(code) => {
                        println!(
                            "Warning: error while querying ChildProcess Audit policy: {:X}",
                            code
                        );
                    }
                }
            }
            Err(-1073741772) => {}
            Err(code) => {
                println!(
                    "Warning: error while querying ChildProcess policy: {:X}",
                    code
                );
            }
        }

        // Query UserShadowStack Policy
        match Ntdll::query_image_mitigation_policy::<RtlImageMitigationUserShadowStackPolicy>(
            exe_name,
            ImageMitigationPolicy::ImageUserShadowStackPolicy,
            0,
        ) {
            Ok(policy) => {
                if let Some(AppPolicyContainer::UserShadowStack(user_shadow_stack_policy)) =
                    mits.policies.get_mut(&PolicyName::UserShadowStack)
                {
                    let mut uss_force = false;
                    let mut strict_mode_option = OptionValue::Off;
                    let uss_value = Self::option_state_to_option_value(
                        policy.user_shadow_stack.policy_state,
                        &mut |force| uss_force = force,
                        &mut |option| strict_mode_option = option,
                    );
                    user_shadow_stack_policy.set_override_user_shadow_stack(uss_force);
                    user_shadow_stack_policy.set_policy(
                        PolicyOptionName::UserShadowStackStrictMode,
                        strict_mode_option,
                    );
                    user_shadow_stack_policy
                        .set_policy(PolicyOptionName::UserShadowStack, uss_value);

                    let ip_val_value = Self::option_state_to_option_value(
                        policy.set_context_ip_validation.policy_state,
                        &mut |_| {},
                        &mut |_| {},
                    );
                    user_shadow_stack_policy
                        .set_policy(PolicyOptionName::SetContextIpValidation, ip_val_value);

                    let cet_value = Self::option_state_to_option_value(
                        policy.block_non_cet_binaries.policy_state,
                        &mut |_| {},
                        &mut |_| {},
                    );
                    user_shadow_stack_policy
                        .set_policy(PolicyOptionName::BlockNonCetBinaries, cet_value);
                }

                // Query audit flags
                match Ntdll::query_image_mitigation_policy::<RtlImageMitigationUserShadowStackPolicy>(
                    exe_name,
                    ImageMitigationPolicy::ImageUserShadowStackPolicy,
                    8,
                ) {
                    Ok(audit_policy) => {
                        if let Some(AppPolicyContainer::UserShadowStack(user_shadow_stack_policy)) =
                            mits.policies.get_mut(&PolicyName::UserShadowStack)
                        {
                            let audit_value = Self::option_state_to_option_value(
                                audit_policy.user_shadow_stack.policy_state,
                                &mut |_| {},
                                &mut |_| {},
                            );
                            user_shadow_stack_policy
                                .set_policy(PolicyOptionName::AuditUserShadowStack, audit_value);
                        }
                    }
                    Err(code) => {
                        println!(
                            "Warning: error while querying UserShadowStack Audit policy: {:X}",
                            code
                        );
                    }
                }
            }
            Err(-1073741772) => {}
            Err(code) => {
                println!(
                    "Warning: error while querying UserShadowStack policy: {:X}",
                    code
                );
            }
        }

        Ok(mits)
    }

    pub fn get_policy_list_from_registry(
        &self,
        process_name: Option<&str>,
    ) -> Result<Vec<AppMitigations>, String> {
        let mut list_from_registry = Vec::new();

        unsafe {
            let mut hkey = HKEY::default();
            let key_path = HSTRING::from(
                "SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\Image File Execution Options",
            );

            if RegOpenKeyExW(HKEY_LOCAL_MACHINE, &key_path, Some(0), KEY_READ, &mut hkey).is_ok() {
                let mut index = 0;
                loop {
                    let mut subkey_name = [0u16; 256];
                    let mut subkey_len = subkey_name.len() as u32;

                    let result = RegEnumKeyExW(
                        hkey,
                        index,
                        Some(windows::core::PWSTR(subkey_name.as_mut_ptr())),
                        &mut subkey_len,
                        None,
                        Some(windows::core::PWSTR::null()),
                        None,
                        None,
                    );

                    if result == ERROR_NO_MORE_ITEMS {
                        break;
                    }

                    if result == ERROR_SUCCESS {
                        let subkey_name_str =
                            String::from_utf16_lossy(&subkey_name[..subkey_len as usize]);

                        if process_name.is_none()
                            || process_name.unwrap().eq_ignore_ascii_case(&subkey_name_str)
                        {
                            let mut subkey = HKEY::default();
                            let subkey_hstring = HSTRING::from(&subkey_name_str);

                            if RegOpenKeyExW(hkey, &subkey_hstring, Some(0), KEY_READ, &mut subkey)
                                .is_ok()
                            {
                                // Check if MitigationOptions or MitigationAuditOptions exist
                                if self.reg_value_exists(subkey, "MitigationOptions")
                                    || self.reg_value_exists(subkey, "MitigationAuditOptions")
                                {
                                    match Self::build_full_policy(Some(&subkey_name_str)) {
                                        Ok(mut app_mitigations) => {
                                            app_mitigations.source = "Registry".to_string();
                                            list_from_registry.push(app_mitigations);
                                        }
                                        Err(error) => {
                                            println!("Warning: {}", error);
                                        }
                                    }
                                }

                                // Check UseFilter
                                if self.reg_value_exists(subkey, "UseFilter") {
                                    let mut filter_index = 0;
                                    loop {
                                        let mut filter_subkey_name = [0u16; 256];
                                        let mut filter_subkey_len = filter_subkey_name.len() as u32;

                                        let filter_result = RegEnumKeyExW(
                                            subkey,
                                            filter_index,
                                            Some(windows::core::PWSTR(
                                                filter_subkey_name.as_mut_ptr(),
                                            )),
                                            &mut filter_subkey_len,
                                            None,
                                            Some(windows::core::PWSTR::null()),
                                            None,
                                            None,
                                        );

                                        if filter_result == ERROR_NO_MORE_ITEMS {
                                            break;
                                        }

                                        if filter_result == ERROR_SUCCESS {
                                            let filter_subkey_name_str = String::from_utf16_lossy(
                                                &filter_subkey_name[..filter_subkey_len as usize],
                                            );
                                            let mut filter_subkey = HKEY::default();
                                            let filter_subkey_hstring =
                                                HSTRING::from(&filter_subkey_name_str);

                                            if RegOpenKeyExW(
                                                subkey,
                                                &filter_subkey_hstring,
                                                Some(0),
                                                KEY_READ,
                                                &mut filter_subkey,
                                            )
                                            .is_ok()
                                            {
                                                if let Some(obj) = self.get_reg_string_value(
                                                    filter_subkey,
                                                    "FilterFullPath",
                                                ) {
                                                    if obj.trim() == "REG_SZ" {
                                                        continue;
                                                    }
                                                    if self.reg_value_exists(
                                                        filter_subkey,
                                                        "MitigationOptions",
                                                    ) || self.reg_value_exists(
                                                        filter_subkey,
                                                        "MitigationAuditOptions",
                                                    ) {
                                                        match Self::build_full_policy(Some(&obj)) {
                                                            Ok(mut app_mitigations) => {
                                                                app_mitigations.source =
                                                                    "Registry".to_string();
                                                                list_from_registry
                                                                    .push(app_mitigations);
                                                            }
                                                            Err(error) => {
                                                                println!("Warning: {}", error);
                                                            }
                                                        }
                                                    }
                                                }
                                            }
                                        }

                                        filter_index += 1;
                                    }
                                }
                            }
                        }
                    }

                    index += 1;
                }
            }
        }

        Ok(list_from_registry)
    }

    pub fn get_policy_from_registry_by_name(
        &self,
        process_name_or_path: &str,
    ) -> Result<Vec<AppMitigations>, String> {
        let file_name = Path::new(process_name_or_path)
            .file_name()
            .and_then(|s| s.to_str())
            .unwrap_or(process_name_or_path);

        let list_from_registry = self.get_policy_list_from_registry(Some(file_name))?;
        let mut from_registry_by_name = Vec::new();
        let flag = file_name != process_name_or_path;

        for app_mitigations in &list_from_registry {
            if (flag
                && app_mitigations
                    .process_name
                    .eq_ignore_ascii_case(process_name_or_path))
                || (!flag
                    && Path::new(&app_mitigations.process_name)
                        .file_name()
                        .and_then(|s| s.to_str())
                        .unwrap_or(&app_mitigations.process_name)
                        .eq_ignore_ascii_case(process_name_or_path))
            {
                from_registry_by_name.push(app_mitigations.clone());
            }
        }

        if flag && from_registry_by_name.is_empty() {
            for app_mitigations in &list_from_registry {
                if app_mitigations.process_name.eq_ignore_ascii_case(file_name) {
                    from_registry_by_name.push(app_mitigations.clone());
                }
            }
        }

        Ok(from_registry_by_name)
    }

    pub fn get_system_policy(&self) -> Result<AppMitigations, String> {
        match Self::build_full_policy(None) {
            Ok(mut system_policy) => {
                system_policy.source = "System Defaults".to_string();
                system_policy.process_name = "System".to_string();
                Ok(system_policy)
            }
            Err(error) => Err(error),
        }
    }

    fn reg_value_exists(&self, hkey: HKEY, value_name: &str) -> bool {
        unsafe {
            let value_name_hstring = HSTRING::from(value_name);
            let mut data_type = REG_VALUE_TYPE::default();
            let mut data_size = 0u32;

            RegQueryValueExW(
                hkey,
                &value_name_hstring,
                None,
                Some(&mut data_type),
                None,
                Some(&mut data_size),
            )
            .is_ok()
        }
    }

    fn get_reg_string_value(&self, hkey: HKEY, value_name: &str) -> Option<String> {
        unsafe {
            let value_name_hstring = HSTRING::from(value_name);
            let mut data_type = REG_VALUE_TYPE::default();
            let mut data_size = 0u32;

            if RegQueryValueExW(
                hkey,
                &value_name_hstring,
                None,
                Some(&mut data_type),
                None,
                Some(&mut data_size),
            )
            .is_ok()
                && data_type == REG_SZ
            {
                let mut buffer = vec![0u16; (data_size / 2) as usize];

                if RegQueryValueExW(
                    hkey,
                    &value_name_hstring,
                    None,
                    Some(&mut data_type),
                    Some(buffer.as_mut_ptr() as *mut u8),
                    Some(&mut data_size),
                )
                .is_ok()
                {
                    let result = String::from_utf16_lossy(&buffer);
                    return Some(result.trim_end_matches('\0').to_string());
                }
            }

            None
        }
    }
}
